import cv2
import mediapipe as mp
import numpy as np

# 初始化 MediaPipe Pose
mp_pose = mp.solutions.pose
mp_drawing = mp.solutions.drawing_utils

# 定义火柴人连接点
stickman_connections = [
    (mp_pose.PoseLandmark.NOSE, mp_pose.PoseLandmark.LEFT_SHOULDER),
    (mp_pose.PoseLandmark.NOSE, mp_pose.PoseLandmark.RIGHT_SHOULDER),
    (mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.RIGHT_SHOULDER),
    (mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.LEFT_ELBOW),
    (mp_pose.PoseLandmark.RIGHT_SHOULDER, mp_pose.PoseLandmark.RIGHT_ELBOW),
    (mp_pose.PoseLandmark.LEFT_ELBOW, mp_pose.PoseLandmark.LEFT_WRIST),
    (mp_pose.PoseLandmark.RIGHT_ELBOW, mp_pose.PoseLandmark.RIGHT_WRIST),
    (mp_pose.PoseLandmark.LEFT_SHOULDER, mp_pose.PoseLandmark.LEFT_HIP),
    (mp_pose.PoseLandmark.RIGHT_SHOULDER, mp_pose.PoseLandmark.RIGHT_HIP),
    (mp_pose.PoseLandmark.LEFT_HIP, mp_pose.PoseLandmark.RIGHT_HIP),
    (mp_pose.PoseLandmark.LEFT_HIP, mp_pose.PoseLandmark.LEFT_KNEE),
    (mp_pose.PoseLandmark.RIGHT_HIP, mp_pose.PoseLandmark.RIGHT_KNEE),
    (mp_pose.PoseLandmark.LEFT_KNEE, mp_pose.PoseLandmark.LEFT_ANKLE),
    (mp_pose.PoseLandmark.RIGHT_KNEE, mp_pose.PoseLandmark.RIGHT_ANKLE),
]

def draw_stickman(image, landmarks, image_width, image_height):
    """绘制火柴人图像"""
    stickman_image = np.ones((image_height, image_width, 3), dtype=np.uint8) * 255  # 创建白色背景
    
    # 绘制连接线
    for connection in stickman_connections:
        start_idx = connection[0]
        end_idx = connection[1]
        
        if (landmarks.landmark[start_idx].visibility > 0.2 and 
            landmarks.landmark[end_idx].visibility > 0.2):
            start_point = (int(landmarks.landmark[start_idx].x * image_width),
                          int(landmarks.landmark[start_idx].y * image_height))
            end_point = (int(landmarks.landmark[end_idx].x * image_width),
                        int(landmarks.landmark[end_idx].y * image_height))
            
            cv2.line(stickman_image, start_point, end_point, (0, 0, 0), 2)
    
    # 绘制关节点
    for idx, landmark in enumerate(landmarks.landmark):
        if landmark.visibility > 0.2:
            x = int(landmark.x * image_width)
            y = int(landmark.y * image_height)
            cv2.circle(stickman_image, (x, y), 5, (0, 0, 255), -1)
    
    return stickman_image

def resize_with_aspect_ratio(image, width=None, height=None, inter=cv2.INTER_AREA):
    """按比例调整图像大小"""
    dim = None
    (h, w) = image.shape[:2]

    if width is None and height is None:
        return image
    if width is None:
        r = height / float(h)
        dim = (int(w * r), height)
    else:
        r = width / float(w)
        dim = (width, int(h * r))

    return cv2.resize(image, dim, interpolation=inter)

def main():
    # 读取本地MP4视频文件
    video_path = 'cxk.mp4'  # 请替换为你的视频文件路径
    cap = cv2.VideoCapture(video_path)
    
    # 检查视频是否成功打开
    if not cap.isOpened():
        print("无法打开视频文件")
        exit()
    
    # 获取视频的宽度、高度和帧率
    original_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    original_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    
    # 设置默认缩放比例
    scale_percent = 40  # 默认缩小到50%
    
    # 初始化 Pose 模型
    with mp_pose.Pose(min_detection_confidence=0.5, min_tracking_confidence=0.5) as pose:
        while True:
            # 读取一帧视频
            ret, frame = cap.read()
            
            # 检查是否成功读取帧
            if not ret:
                print("无法获取帧或视频已结束")
                break
            
            # 计算调整后的尺寸
            width = int(original_width * scale_percent / 100)
            height = int(original_height * scale_percent / 100)
            
            # 调整帧大小
            frame = resize_with_aspect_ratio(frame, width=width)
            
            # 转换为 RGB 格式
            image_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            
            # 处理图像，检测姿势
            results = pose.process(image_rgb)
            
            # 在原始图像上绘制关节点和连接线
            annotated_image = frame.copy()
            if results.pose_landmarks:
                mp_drawing.draw_landmarks(
                    annotated_image,
                    results.pose_landmarks,
                    mp_pose.POSE_CONNECTIONS,
                    mp_drawing.DrawingSpec(color=(0, 255, 0), thickness=2, circle_radius=2),
                    mp_drawing.DrawingSpec(color=(0, 0, 255), thickness=2)
                )
                
                # 绘制火柴人
                stickman_image = draw_stickman(annotated_image, results.pose_landmarks, width, height)
                
                # 显示火柴人窗口
                cv2.imshow('Stickman', stickman_image)
            
            # 显示控制提示
            cv2.putText(annotated_image, f'Scale: {scale_percent}%', (10, 30), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            cv2.putText(annotated_image, 'Press + to zoom in, - to zoom out', (10, 60), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
            cv2.putText(annotated_image, 'Press SPACE to pause, Q to quit', (10, 85), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
            
            # 显示原始图像窗口
            cv2.imshow('Pose Tracking', annotated_image)
            
            # 控制视频播放速度
            wait_time = int(200/fps)  # 根据视频帧率计算等待时间
            
            # 按键处理
            key = cv2.waitKey(wait_time) & 0xFF
            
            if key == ord('q'):  # 退出程序
                break
            elif key == ord(' '):  # 空格键暂停/继续
                cv2.waitKey(0)  # 等待任意按键继续
            elif key == ord('+') or key == ord('='):  # 放大
                scale_percent = min(scale_percent + 10, 200)  # 最大放大到200%
            elif key == ord('-'):  # 缩小
                scale_percent = max(scale_percent - 10, 10)  # 最小缩小到10%
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()